(* Mathematica Package *)

BeginPackage["FewSpinDynamics`", { "MathMR`", "MRData`", "NVInteraction`", "Hamiltonian`"}]
(* Exported symbols added here with SymbolName::usage *)  
ProjectionOp::usage="";
EspinFlipPI::usage=""; 
EspinFlipPIhalf::usage="";
EspinPopulation::usage="";
EspinPolarization::usage="";

HamiltonianRotatingFrame::usage="Return the Hamiltonian in rotating frame";


Begin["`Private`"] (* Begin Private Context *) 

HamiltonianRotatingFrame[nvdir_, b_, theta_, phi_, nuc_, mwfreq_, rabifreq_, mwphase_] := 
	Module[{res, idmat, nvhami, mwHami, mwHami1, rabiHami, rabiHami1},
		idmat=If[nuc=="N14",IdentityMatrix[3],IdentityMatrix[2]]; 
		nvhami=DiagonalMatrix[Diagonal[SingleNV[nvdir, b, theta, phi, nuc]]];
		
		mwHami=2.*Pi*1000*{{mwfreq,0,0},{0,0,0},{0,0,mwfreq}};(*MHz ---> k rad/s*)
		mwHami1= KroneckerProduct[mwHami, idmat];
		
		rabiHami=2.*Pi*1000*rabifreq*( Cos[mwphase]*SPIN1x + Sin[mwphase]*SPIN1y );
		rabiHami1= KroneckerProduct[rabiHami,idmat];
		
		res=nvhami - mwHami1 + rabiHami1;
		Return[res];
	]

ProjectionOp[fromState_,toState_]:= KroneckerProduct[toState,Conjugate[fromState] ]

EspinFlipPI[nvdir_, state1_, state2_]:=
	Module[{res, unChangeState, stateVect1, stateVect2, unChangeVect},
		unChangeState = Complement[{1,2,3},{state1,state2}][[1]];
		stateVect1 = EigenState[nvdir, state1];
		stateVect2 = EigenState[nvdir, state2];
		unChangeVect = EigenState[nvdir, unChangeState];
		
		res = -I * ProjectionOp[stateVect1, stateVect2]
			-I * ProjectionOp[stateVect2, stateVect1]
			+ ProjectionOp[unChangeVect, unChangeVect];
		Return[res];
	]
	
EspinFlipPIhalf[nvdir_, state1_, state2_]:=
	Module[{res, unChangeState, stateVect1, stateVect2, unChangeVect},
		unChangeState = Complement[{1,2,3},{state1,state2}][[1]];
		stateVect1 = EigenState[nvdir, state1];
		stateVect2 = EigenState[nvdir, state2];
		unChangeVect = EigenState[nvdir, unChangeState];
		
		res = 1/Sqrt[2] * 
			( ProjectionOp[stateVect1, stateVect1]
			- I * ProjectionOp[stateVect1, stateVect2]
			- I * ProjectionOp[stateVect2, stateVect1]
			+ ProjectionOp[stateVect2, stateVect2]
			) + ProjectionOp[unChangeVect, unChangeVect];
		Return[res];
	]

EspinPopulation[nvdir_, state_]:=
	Module[{res, stateVect},
		stateVect = EigenState[nvdir, state];
		res = ProjectionOp[stateVect, stateVect];
		Return[res]
	]
	
EspinPolarization[nvdir_, state1_, state2_]:=
	Module[{res, stateVect1, stateVect2},
		stateVect1 = EigenState[nvdir, state1];
		stateVect2 = EigenState[nvdir, state2];
		res = ProjectionOp[stateVect1, stateVect2];
		Return[res]
	]
	
End[] (* End Private Context *)

EndPackage[]